# -*-coding: utf-8 -*-
"""
    代码需要用到pybaseutils工具，请使用pip安装即可：pip install pybaseutils
"""
 
import os
import cv2
import numpy as np
from tqdm import tqdm
from pybaseutils import file_utils, image_utils


def cv_imwrite(file_path, img, params=None):
    """
    支持中文路径的图片保存函数
    :param file_path: 包含中文的完整保存路径
    :param img: OpenCV图像矩阵
    :param params: 编码参数 如 [cv2.IMWRITE_JPEG_QUALITY, 90]
    """
    try:
        # 自动创建目录
        dir_path = os.path.dirname(file_path)
        if dir_path and not os.path.exists(dir_path):
            os.makedirs(dir_path, exist_ok=True)
            
        # 获取文件后缀
        ext = os.path.splitext(file_path)[1]
        
        # 使用imencode进行编码
        ret, buf = cv2.imencode(ext, img, params)
        
        if ret:
            with open(file_path, 'wb') as f:
                f.write(buf.tobytes())
            return True
        return False
    except Exception as e:
        print(f'Save image failed: {e}')
        return False
 
def get_plate_licenses(plate):
    """
    普通蓝牌共有7位字符；新能源车牌有8位字符： https://baike.baidu.com/item/%E8%BD%A6%E7%89%8C/8347320?fr=aladdin
    《新能源电动汽车牌照和普通牌照区别介绍》https://www.yoojia.com/ask/4-11906976349117851507.html
    新能源汽车车牌可分为三部分：省份简称(1位汉字)十地方行政区代号(1位字母)十序号(6位)
    字母“D”代表纯电动汽车；
    字母“F”代表非纯电动汽车(包括插电式混合动力和燃料电池汽车等)。
    :param plate:
    :return:
    """
    provinces = ["皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤",
                 "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "警", "学", "O"]
    alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V',
                 'W', 'X', 'Y', 'Z', 'O']
    ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
           'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O']
    result = [provinces[int(plate[0])], alphabets[int(plate[1])]]
    result += [ads[int(p)] for p in plate[2:]]
    result = "".join(result)
    # 新能源车牌的要求，如果不是新能源车牌可以删掉这个if
    # if result[2] != 'D' and result[2] != 'F' \
    #         and result[-1] != 'D' and result[-1] != 'F':
    #     print(plate)
    #     print("Error label, Please check!")
    print(plate, result)
    return result
 
 
def parser_annotations(image_file):
    """
    :param image_file: 图片路径
    :return: 返回标注信息info
    """
    filename = os.path.basename(image_file)
    try:
        annotations = filename.split("-")
        rate = annotations[0]  # 车牌区域占整个画面的比例；
        angle = annotations[1].split("_")  # 车牌水平和垂直角度, 水平95°, 竖直113°
        box = annotations[2].replace("&", "_").split("_")  # 标注框左上、右下坐标，左上(154, 383), 右下(386, 473)
        point = annotations[3].replace("&", "_").split("_")  # 标注框四个角点坐标，顺序为右下、左下、左上、右上
        plate = annotations[4].split("_")  # licenses 标注框四个角点坐标，顺序为右下、左下、左上、右上
        plate = get_plate_licenses(plate)
        box = [int(b) for b in box]
        point = [int(b) for b in point]
        point = np.asarray(point).reshape(-1, 2)
        bboxes = [box]
        angles = [angle]
        points = [point]
        plates = [plate]
        labels = ["plate"] * len(bboxes)
    except Exception as e:
        bboxes = []
        points = []
        labels = []
        plates = []
        angles = []
    info = {"filename": filename, "bboxes": bboxes, "points": points,
            "labels": labels, "plates": plates, "angles": angles}
    return info
 
 
def save_plate_licenses(image, bboxes, plates, out_dir, name=""):
    crops = image_utils.get_bboxes_crop(image, bboxes)
    for i in range(len(crops)):
        label = plates[i]
        # image_id = file_utils.get_time(format="p")
        file = os.path.join(out_dir, "{}_{}_{:0=3d}.jpg".format(label, name, i))
        file_utils.create_file_path(file)
        cv_imwrite(file, crops[i])
 
 
def converter_CCPD2yolo(image_dir, output_dir, vis=False):
    """
    将CCPD数据集转换为YOLO格式
    :param image_dir: CCPD数据集图片根目录
    :param output_dir: 输出YOLO格式数据集目录
    :param vis: 是否可视化效果
    """
    print("image_dir: {}".format(image_dir))
    print("output_dir: {}".format(output_dir))
    
    # 创建输出目录结构
    images_dir = os.path.join(output_dir, "images")
    labels_dir = os.path.join(output_dir, "labels")
    file_utils.create_dir(images_dir)
    file_utils.create_dir(labels_dir)
    
    class_set = []
    image_list = file_utils.get_images_list(image_dir)
    valid_count = 0
    
    for i, image_file in enumerate(tqdm(image_list)):
        info = parser_annotations(image_file)
        labels = info["labels"]
        bboxes = info["bboxes"]
        points = info["points"]
        plates = info["plates"]
        angles = info["angles"]
        image_name = info["filename"]
        
        if len(labels) == 0:
            continue
            
        if not os.path.exists(image_file):
            print("not exist: {}".format(image_file))
            continue
            
        # 处理中文路径问题，使用np.fromfile和cv2.imdecode
        try:
            # 使用np.fromfile读取图片文件，支持中文路径
            img_array = np.fromfile(image_file, dtype=np.uint8)
            # 使用cv2.imdecode解码图片
            image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            if image is None:
                print("Failed to decode image: {}".format(image_file))
                continue
        except Exception as e:
            print("Failed to load image: {}, Error: {}".format(image_file, str(e)))
            continue
            
        # 获取图像尺寸
        img_height, img_width = image.shape[:2]
        
        # 保存图像
        image_name = os.path.basename(image_name)
        img_postfix = image_name.split(".")[-1]
        image_id = image_name[:-len(img_postfix) - 1]
        
        # 新的文件名（简化文件名，避免过长）
        new_image_name = "{:06d}.jpg".format(valid_count)
        new_image_path = os.path.join(images_dir, new_image_name)
        
        # 保存图像文件
        cv_imwrite(new_image_path, image)
        
        # 创建YOLO格式标签文件
        label_name = "{:06d}.txt".format(valid_count)
        label_path = os.path.join(labels_dir, label_name)
        
        with open(label_path, 'w', encoding='utf-8') as f:
            # 同时遍历bboxes和points
            for j, (label, bbox, point) in enumerate(zip(labels, bboxes, points)):
                if label=="plate":
                    class_id = 0
                xmin, ymin, xmax, ymax = bbox
                
                # 转换bbox坐标
                center_x = (xmin + xmax) / 2.0 / img_width
                center_y = (ymin + ymax) / 2.0 / img_height
                width = (xmax - xmin) / img_width
                height = (ymax - ymin) / img_height
                
                # 转换关键点坐标（示例point格式：[[x1,y1],...[x4,y4]]）
                kps_normalized = []
                for (x, y) in point:
                    kps_normalized.extend([
                        x/img_width,   # 归一化x坐标
                        y/img_height  # 归一化y坐标
                        # 1              # 可见性参数
                    ])
                
                # 构建完整标签行（5+12=17列），由于不需要关键点可视化列，修改一下
                label_line = [
                    class_id,
                    center_x, center_y, width, height,
                    *kps_normalized
                ]
                
                # 格式化为字符串
                line = ' '.join(f'{x:.6f}' if isinstance(x, float) else str(x) 
                              for x in label_line)
                f.write(f'{line}\n')
        
        class_set = labels + class_set
        class_set = list(set(class_set))
        valid_count += 1
        
        if vis and valid_count <= 10:  # 只显示前10张图片
            vis_image = image.copy()
            vis_image = image_utils.draw_image_bboxes_text(vis_image, bboxes, plates, 
                                                         color=(255, 0, 0), thickness=3,
                                                         fontScale=1.2, drawType="chinese")
            image_utils.cv_show_image("det", vis_image, use_rgb=False, delay=1000)
            
        # 每处理1000张图片打印一次进度
        if valid_count % 1000 == 0:
            print("Processed {} valid images".format(valid_count))
    
    print("Total processed {} valid images".format(valid_count))
    print("class_set: {}".format(class_set))
    print("Images saved to: {}".format(images_dir))
    print("Labels saved to: {}".format(labels_dir))
 
 
if __name__ == "__main__":
    image_dir = r"D:\songlin\data\9.公开数据集\CCPD2020\ccpd_green\images\val"
    output_dir = r"D:\songlin\data\9.公开数据集\CCPD2020\ccpd_green\images\val_labels"
    converter_CCPD2yolo(image_dir, output_dir,vis=False)